package com.ruoyi.framework.aspectj.commonservice;

import org.springframework.util.StringUtils;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.UUID;

/**
 * @author: silver
 * created:11/21/22
 * description
 */
public class TraceIdHandlerInterceptor extends HandlerInterceptorAdapter {
    static final String REQUEST_ID_QUERY_KEY = "_reqid";
    static final String REQUEST_ID_HEADER_KEY = "x-api-reqid";

    public TraceIdHandlerInterceptor() {
    }

    public boolean preHandle(final HttpServletRequest request, final HttpServletResponse response, final Object handler) {
        String requestId = request.getParameter(REQUEST_ID_QUERY_KEY);
        if (StringUtils.isEmpty(requestId)) {
            requestId = request.getHeader(REQUEST_ID_HEADER_KEY);
        }

        if (StringUtils.isEmpty(requestId)) {
            requestId = UUID.randomUUID().toString();
        }

        TraceIdContext.set(requestId);
        response.setHeader(REQUEST_ID_HEADER_KEY, requestId);
        return true;
    }

    public void afterCompletion(final HttpServletRequest request, final HttpServletResponse response, final Object handler, final Exception ex) {
        TraceIdContext.remove();
    }
}
